package com.bsfcs.rag.service;

import com.bsfcs.rag.entity.ChatParam;
import com.bsfcs.rag.repository.TenantAwareChatMemoryRepository;
import com.bsfcs.rag.utils.FilePartUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.Resource;
import org.springframework.http.codec.multipart.FilePart;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

import org.springframework.ai.openai.OpenAiEmbeddingModel;

@Service
public class DocumentIngestionService {

    private static final Logger logger = LoggerFactory.getLogger(DocumentIngestionService.class);

    private final ChatClient chatClient;

    private final ChatMemory chatMemory;

    private final VectorStore vectorStore;

    private final ChatModel chatModel;

    private final OpenAiEmbeddingModel embeddingClient;

    // Chunk上下文提示模板
    private static final String CHUNK_CONTEXT_PROMPT = """
            这是文档中的一个片段
            <chunk>
            {chunk_content}
            </chunk>
            
            请生成简洁的上下文描述，以便提升该片段在全文检索和语义理解中的效果。
            只回答简洁的上下文，不要回答其他问题。
            """;

    // 反向HyDE提示模板 - 为chunk生成相关问题
    private static final String REVERSE_HYDE_PROMPT = """
            给定以下文本块，请生成 3 个不同的问题，使该文本块能够作为理想答案进行匹配：
            
            Chunk: {chunk_content}
            
            问题（用1.、2.等列举问题）：
            """;

    public DocumentIngestionService(ChatClient.Builder chatClientBuilder,
                                    VectorStore vectorStore,
                                    ChatModel chatModel,
                                    OpenAiEmbeddingModel embeddingClient,
                                    TenantAwareChatMemoryRepository chatMemoryRepository) {
        this.vectorStore = vectorStore;
        this.chatModel = chatModel;
        this.embeddingClient = embeddingClient;
        this.chatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(chatMemoryRepository)
                .maxMessages(10)
                .build();

        this.chatClient = chatClientBuilder.build();
    }

    public Mono<Void> ingestDocuments(String knowledgeBaseId,
                                      String knowledgeBaseName,
                                      String fileArchiveId,
                                      String fileId,
                                      FilePart file) {

        return FilePartUtils.toResource(file)
                .flatMap(resource -> {
                    if (!resource.exists()) {
                        return Mono.error(new IllegalArgumentException("Uploaded file does not exist"));
                    }
                    return processResource(resource, knowledgeBaseId, knowledgeBaseName, fileArchiveId, fileId);
                });
    }

    /**
     * 统一处理Resource的文档向量化流程
     *
     * @param resource          资源文件
     * @param knowledgeBaseId   知识库ID
     * @param knowledgeBaseName 知识库名称
     * @param fileArchiveId     文件归档栏目ID
     * @param fileId            文件ID（可选）
     */
    private Mono<Void> processResource(Resource resource, String knowledgeBaseId, String knowledgeBaseName, String fileArchiveId, String fileId) {
        return Mono.fromRunnable(() -> {
            try {
                logger.info("处理文件: {}", resource.getFilename());

                // 1. Extract (提取): 使用 TikaDocumentReader 加载文档
                TikaDocumentReader documentReader = new TikaDocumentReader(resource);
                List<Document> documents = documentReader.get();
                logger.info("成功加载文档，共 {} 个部分。", documents.size());

                // 为每个文档添加知识库元数据
                for (Document doc : documents) {
                    Map<String, Object> metadata = doc.getMetadata();
                    metadata.put("knowledge_base_id", knowledgeBaseId);
                    metadata.put("knowledge_base_name", knowledgeBaseName);
                    metadata.put("file_archive_id", fileArchiveId);
                    metadata.put("filename", resource.getFilename());
                    if (fileId != null) {
                        metadata.put("fileId", fileId);
                    }
                }

                // 获取完整文档内容用于上下文理解
                StringBuilder fullDocumentContent = new StringBuilder();
                for (Document doc : documents) {
                    fullDocumentContent.append(doc.getText()).append("\n");
                }

                // 2. Transform (转换): 使用 TokenTextSplitter 进行语义切分
                ChineseTextSplitter chineseTextSplitter = new ChineseTextSplitter();

                List<Document> splitDocuments = chineseTextSplitter.splitChineseDocument(fullDocumentContent.toString());

                // 为切分后的文档块添加知识库元数据
                for (Document doc : splitDocuments) {
                    Map<String, Object> metadata = doc.getMetadata();
                    metadata.put("knowledge_base_id", knowledgeBaseId);
                    metadata.put("knowledge_base_name", knowledgeBaseName);
                    metadata.put("file_archive_id", fileArchiveId);
                    metadata.put("filename", resource.getFilename());
                }

                logger.info("文档切分完成，共得到 {} 个文本块。", splitDocuments.size());

                // 3. 上下文和反向HyDE增强: 为每个chunk生成上下文摘要和相关问题
                logger.info("开始进行上下文和反向HyDE增强...");
                enhanceDocuments(splitDocuments)
                        .subscribe(enhancedDocuments -> {
                            // 4. Load (加载): 将增强后的文档块添加到 VectorStore 中
                            // 分批处理，避免超过DashScope的批处理大小限制(最大10)
                            int batchSize = 10;
                            for (int i = 0; i < enhancedDocuments.size(); i += batchSize) {
                                int endIndex = Math.min(i + batchSize, enhancedDocuments.size());
                                List<Document> batch = enhancedDocuments.subList(i, endIndex);
                                vectorStore.delete(batch.stream().map(Document::getId).toList());
                                vectorStore.add(batch);
                            }
                            logger.info("文档向量化并存入 VectorStore 成功！");
                        }, error -> {
                            logger.error("处理文件 {} 时出错: {}", resource.getFilename(), error.getMessage(), error);
                        });
            } catch (Exception e) {
                logger.error("处理文件 {} 时出错: {}", resource.getFilename(), e.getMessage(), e);
            }
        }).then();
    }


    /**
     * 增强单个文档，结合上下文增强和反向HyDE增强
     *
     * @param documents 需要增强的文档
     * @return 增强后的文档
     */
    private Mono<List<Document>> enhanceDocuments(List<Document> documents) {
        return Mono.fromCallable(() -> {
            return documents.stream().map(document -> {
                try {
                    String chunkContent = document.getText();
                    StringBuilder enhancedContent = new StringBuilder();

                    // 1. 添加上下文信息
                    try {
                        PromptTemplate contextPromptTemplate = new PromptTemplate(CHUNK_CONTEXT_PROMPT);
                        Prompt contextPrompt = contextPromptTemplate.create(Map.of(
                                "chunk_content", chunkContent
                        ));
                        String context = chatModel.call(contextPrompt).getResult().getOutput().getText();
                        enhancedContent.append("Context: ").append(context).append("\n\n");
                    } catch (Exception e) {
                        logger.warn("为文档生成上下文时出错: {}", e.getMessage());
                    }

                    // 2. 添加反向HyDE问题
                    try {
                        PromptTemplate hydePromptTemplate = new PromptTemplate(REVERSE_HYDE_PROMPT);
                        Prompt hydePrompt = hydePromptTemplate.create(Map.of(
                                "chunk_content", chunkContent
                        ));
                        String questions = chatModel.call(hydePrompt).getResult().getOutput().getText();
                        enhancedContent.append("Related Questions:\n").append(questions).append("\n\n");
                    } catch (Exception e) {
                        logger.warn("为文档生成反向HyDE问题时出错: {}", e.getMessage());
                    }

                    // 3. 添加原始内容
                    enhancedContent.append("Content:\n").append(chunkContent);

                    // 创建新的增强文档
                    return new Document(enhancedContent.toString(), document.getMetadata());
                } catch (Exception e) {
                    logger.warn("为文档生成增强内容时出错: {}", e.getMessage());
                    // 出错时返回原始文档
                    return document;
                }
            }).toList();
        }).subscribeOn(Schedulers.boundedElastic());
    }

    /**
     * 根据相似度查询相关文档
     *
     * @param chatParam 查询参数
     * @return 相似度最高的文档列表
     */
    public Flux<String> similaritySearch(ChatParam chatParam) {
        int topK = 2;
        logger.info("开始执行相似度查询: {} (返回数量: {})", chatParam.getUserInput(), topK);

        return Mono.fromCallable(() -> {
                    StringBuilder content = new StringBuilder();
                    try {
                        // 执行普通的相似度查询
                        List<Document> results = vectorStore.similaritySearch(chatParam.getUserInput());

                        // 如果需要限制返回结果数量，可以进行截取
                        if (results.size() > topK) {
                            results = results.subList(0, topK);
                        }

                        logger.info("相似度查询完成，返回 {} 个结果", results.size());

                        results.forEach(i -> {
                            content.append(i.getText()).append("\n\n");
                        });
                    } catch (Exception e) {
                        logger.error("执行相似度查询时出错: {}", e.getMessage(), e);
                    }
                    return content.toString();
                })
                .subscribeOn(Schedulers.boundedElastic())
                .flatMapMany(content -> {
                    if (content.isEmpty()) {
                        return Flux.just("content为空，未找到相关信息");
                    }
                    return this.chatClient.prompt()
                            .advisors(MessageChatMemoryAdvisor.builder(chatMemory).build())

                            .advisors(advisor -> advisor.param(ChatMemory.CONVERSATION_ID,
                                    chatParam.getConversationId() + ":" + chatParam.getUserId() + ":" + chatParam.getTenantId()))
                            .system(systemSpec -> systemSpec
                                    .text("""
                                        你是一名严谨的问答助手。
                                          - 你必须严格依据提供的检索资料回答。
                                          - 不允许编造或补充资料中不存在的信息。
                                          - 如果资料不足以回答，请直接回答：“未找到相关信息”。
                                          - 回答要简洁、准确。
                                      """)
                                    .param("context", content))
                            .user(userSpec -> userSpec
                                    .text("""
                                        请根据以下文档内容回答问题:
                                        
                                        {context}
                                        
                                        用户问题: {question}
                                        """)
                                    .param("context", content)
                                    .param("question", chatParam.getUserInput()))
                            .stream()
                            .content();
                });
    }


    /**
     * 主流程：Embedding 校验 -> LLM 自校验 -> 输出
     */
    private String strictValidationAnswer(String answer, List<String> retrievedChunks) {
        // 1. Embedding 逐句校验
        boolean embeddingPass = embeddingValidation(answer, retrievedChunks);
        logger.info("Embedding 校验结果：{}", embeddingPass ? "通过" : "未通过");
        if (!embeddingPass) {
            return "未找到相关信息";
        }

        // 2. LLM 自校验
        boolean llmPass = llmValidation(answer, retrievedChunks);
        logger.info("LLM 自校验结果：{}", llmPass ? "通过" : "未通过");
        if (!llmPass) {
            return "未找到相关信息";
        }

        // 3. 两道保险都通过
        return answer;
    }

    /**
     * LLM 自校验
     */
    public boolean llmValidation(String answer, List<String> retrievedChunks) {
        String validationPrompt = "以下是检索到的资料:\n" +
                String.join("\n", retrievedChunks) +
                "\n\n以下是生成的答案:\n" + answer +
                "\n\n问题: 该答案是否完全基于检索资料，没有引入资料外的信息？只回答 是 或 否。";

        String result = chatModel.call(new Prompt(validationPrompt))
                .getResult().getOutput().getText().trim();

        boolean pass = result.equals("是");
        logger.info("LLM 校验判定：{} | 模型返回：{}", pass ? "通过" : "未通过", result);
        return pass;
    }

    /**
     * 简单句子分割
     */
    public List<String> splitIntoSentences(String text) {
        String[] sentences = text.split("(?<=[。！？.!?])\\s*");
        return Arrays.stream(sentences)
                .map(String::trim)
                .filter(s -> !s.isEmpty())
                .toList();
    }

    /**
     * 余弦相似度
     */
    private double cosineSimilarity(float[] a, float[] b) {
        double dot = 0.0, normA = 0.0, normB = 0.0;
        for (int i = 0; i < a.length; i++) {
            dot += a[i] * b[i];
            normA += a[i] * a[i];
            normB += b[i] * b[i];
        }
        return dot / (Math.sqrt(normA) * Math.sqrt(normB));
    }

    /**
     * 语义相似度校验（逐句检测 + 日志输出）
     */
    public boolean embeddingValidation(String answer, List<String> retrievedChunks) {
        List<String> sentences = splitIntoSentences(answer);
        if (sentences.isEmpty()) {
            logger.warn("答案为空，判定为不合格");
            return false;
        }

        // 预先计算 chunks embedding
        List<float[]> chunkVecs = embeddingClient.embed(retrievedChunks);

        for (String sentence : sentences) {
            float[] ansVec = embeddingClient.embed(sentence);


            double maxSim = chunkVecs.stream()
                    .mapToDouble(vec -> cosineSimilarity(ansVec, vec))
                    .max().orElse(0.0);

            if (maxSim < 0.7) {
                logger.warn("句子校验未通过：\"{}\" | 最大相似度={} < 0.7", sentence, maxSim);
                return false;
            } else {
                logger.info("句子校验通过：\"{}\" | 最大相似度={}", sentence, maxSim);
            }
        }
        return true;
    }
}